{-# LANGUAGE FlexibleContexts #-}

module Zeno.Core
(
  module Zeno.Clause,
  module Zeno.Type,
  module Zeno.Term,
  module Zeno.FunctionDef,
  ZVar (..), ZVarClass (..), ZDataType (..), Zeno (..),
  ZId, ZTerm, ZType, ZClause, ZEquality, ZProof,
  ZFunctionDef, ZExpression, ZWhen, ZProofStep,
  ZenoState, ZenoReader, ZTermSubstitution,
  instantiate, applyTermDefinitions, simpleTypeTerm,
  newIdS, christen, largestId, getName, emptyZeno, getGlobalVar,
  addLemma, addFunctionDef, addDataType, addGlobalVar,
  getLemma, getFunctionDef, applyFunctionDef, freshFunctionDef,
  findDataType, termType, recursiveDataType,
  functionTerm, variableTerm, constructorTerm, functionTypeTerm,
  showZClause, showZTerm, showZEquality, showZVar, showZId, showZTermSub,
  showZProof, showZExpression, showZFunctionDef, variables,
  irreducibleTerm, groundTerm, unequalTerms, freshVariable
)
where

import Prelude ()
import StdImports
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IMap
import qualified Data.Map as Map
import qualified Data.Set as Set

import Zeno.Clause
import Zeno.Type
import Zeno.Term
import Zeno.FunctionDef
import Zeno.ProofLogic

type ZId = Int
type ZTerm = Term ZVar
type ZType = Type ZDataType
type ZClause = Clause ZVar
type ZEquality = Equality ZVar
type ZTermSubstitution = TermSubstitution ZVar
type ZProofStep = ProofStep ZVar
type ZProof = Proof ZProofStep
type ZFunctionDef = FunctionDef ZVar
type ZExpression = Expression ZVar
type ZWhen = When ZVar

data ZVar
  = ZVar        { zVarId :: !ZId,
                  zVarType :: !ZType,
                  zVarClass :: !ZVarClass }
                  
instance Eq ZVar where
  var1 == var2 = zVarId var1 == zVarId var2
  
instance Ord ZVar where
  var1 `compare` var2 = 
    zVarId var1 `compare` zVarId var2

data ZVarClass 
  = ZVariable 
  | ZFunction 
  | ZConstructor
  deriving ( Eq )

data ZDataType 
  = ZDataType   { zDataTypeId :: !ZId,
                  zDataTypeCons :: [ZVar] }

instance Eq ZDataType where
  dt1 == dt2 = zDataTypeId dt1 == zDataTypeId dt2
  
data Zeno
  = Zeno        { zFunctions :: !(IntMap ZFunctionDef),
                  zGlobalVars :: !(Map String ZVar),
                  zLemmas :: !(Map String ZClause),
                  zDataTypes :: ![ZDataType],
                  zNames :: Map ZId String,
                  zCounter :: !ZId }
                  
instance IdCounter Zeno where
  newId zeno = 
    let id = zCounter zeno + 1
    in (id, zeno { zCounter = id })
                  
emptyZeno 
  = Zeno        { zFunctions = mempty,
                  zGlobalVars = mempty,
                  zLemmas = mempty,
                  zDataTypes = mempty,
                  zNames = mempty,
                  zCounter = 1 }
                  
type ZenoState = State Zeno
type ZenoReader = Reader Zeno

instance NFData Zeno where
  rnf (Zeno funs vars lems dts names ids) =
    rnf funs `seq` rnf vars `seq` 
      rnf lems `seq` rnf dts `seq` rnf names `seq` rnf ids
  
instance NFData ZVar where
  rnf (ZVar id t cls) = rnf id `seq` rnf t
  
instance NFData ZDataType where
  rnf (ZDataType id cons) = rnf id

functionTerm :: ZTerm -> Bool
functionTerm term = 
  zVarClass (termFunction term) == ZFunction

irreducibleTerm :: ZTerm -> Bool
irreducibleTerm = all ((/= ZFunction) . zVarClass)

groundTerm :: ZTerm -> Bool
groundTerm = all ((== ZConstructor) . zVarClass)

unequalTerms :: ZTerm -> ZTerm -> Bool
unequalTerms t1 t2 = 
  constructorTerm t1 &&
  constructorTerm t2 &&
  termFunction t1 /= termFunction t2
  
variables :: Foldable f => f ZVar -> [ZVar]
variables = nub . filter ((== ZVariable) . zVarClass) . toList

variableTerm :: ZTerm -> Bool
variableTerm term = 
  zVarClass (termFunction term) == ZVariable

constructorTerm :: ZTerm -> Bool
constructorTerm term = 
  zVarClass (termFunction term) == ZConstructor
  
getLemma :: MonadReader Zeno m => String -> m ZClause
getLemma name = do
  lemmas <- zLemmas `liftM` ask
  case Map.lookup name lemmas of
    Nothing -> error $ "Lemma \"" ++ name ++ "\" not found"
    Just lemma -> return lemma
    
getFunctionDef :: MonadReader Zeno m => ZVar -> m ZFunctionDef
getFunctionDef var = do
  defs <- zFunctions `liftM` ask
  case IMap.lookup (zVarId var) defs of
    Just def -> return def
    
freshVariable :: ZVar -> ZenoState ZVar
freshVariable (ZVar id typ cls) = do
  new_id <- newIdS
  return $ ZVar new_id typ cls
    
freshFunctionDef :: ZVar -> ZenoState ZFunctionDef
freshFunctionDef var = do
  func_def <- readonly $ getFunctionDef var
  let all_vars = filter ((== ZVariable) . zVarClass) $ 
        toList func_def
  new_vars <- mapM freshVariable all_vars
  let mapping = all_vars `zip` new_vars
  return $ replaceMany mapping func_def

-- This function may give an infinite loop for 
--   non-primitive recursive function definitions
applyFunctionDef :: [ZEquality] -> ZTerm -> ZenoReader (Maybe ZExpression)
applyFunctionDef defs term = do
  let (Var func_var) : args = flattenTerm term
  if zVarClass func_var /= ZFunction || null args
    then return Nothing
    else do
      func_def <- getFunctionDef func_var
      let params = map Var $ funcParameters func_def
      if length params /= length args
        then return Nothing
        else do
          let arg_mapping = Set.fromList $ params `zip` args
              expr = funcExpression func_def 
              mapped_expr = applyTermSubstitution arg_mapping expr
          Just <$> simplifyExpression mapped_expr
  where
    simplifyExpression :: ZExpression -> ZenoReader ZExpression
    simplifyExpression expr@(Value term) = return expr              
    simplifyExpression expr@(CaseOf term whens) = do  
      let ((Var func) : args) = flattenTerm term
          
          matchWhen (When term' expr') = do
            let ((Var func') : args') = flattenTerm term'
            guard (func' == func) 
            let arg_mapping = Set.fromList $ args' `zip` args
            return $ applyTermSubstitution arg_mapping expr'
      
      simpl <- case mapMaybe matchWhen whens of
        [expr'] -> simplifyExpression expr'
        [] -> do
          maybe_expr <- applyFunctionDef defs term
          case maybe_expr of
            Just (Value term') -> do
              let expr' = CaseOf term' whens
              if constructorTerm term 
                then simplifyExpression expr'
                else return expr'
            _ -> return expr
      
      let expr' = applyTermSubstitution (Set.fromList defs) simpl
      if expr' /= expr
        then simplifyExpression expr'
        else return expr
            
recursiveDataType :: ZDataType -> Bool
recursiveDataType datatype = 
  any (any (== (Type datatype)) . dropEnd . flattenType . zVarType) 
    (zDataTypeCons datatype)
  where dropEnd :: [a] -> [a]
        dropEnd ls = take (length ls - 1) ls
  
validClause :: ZClause -> Bool
validClause (Clause vars head body) = all validBodyEq body
  where validBodyEq :: ZEquality -> Bool
        validBodyEq (left_term, right_term) = 
          functionTerm left_term && not (functionTerm right_term)

termType :: ZTerm -> ZType
termType term = 
  let ((Var func) : args) = flattenTerm term
      arg_types = flattenType (zVarType func)
      res_types = drop (length args) arg_types 
  in unflattenType res_types

functionTypeTerm :: ZTerm -> Bool  
functionTypeTerm = functionType . termType

simpleTypeTerm :: ZTerm -> Bool
simpleTypeTerm = simpleType . termType

christen :: MonadState Zeno m => ZId -> String -> m ()
christen id name = do
  zeno <- get
  let names = Map.insert id name (zNames zeno)
  put $ zeno { zNames = names }
  
addDataType :: MonadState Zeno m => ZDataType -> m ()
addDataType dtype = dtype `seq`
  modify $ \z -> z { zDataTypes = dtype : (zDataTypes z) }
  {-
addFunctions :: MonadState Zeno m => [ZVar] -> m ()
addFunctions funcs =
  modify $ \z -> z { zFunctions = funcs ++ (zFunctions z) }
  
addVariables :: MonadState Zeno m => [ZVar] -> m ()
addVariables vars =
  modify $ \z -> z { zVariables = vars ++ (zVariables z) }
  -}
addFunctionDef :: MonadState Zeno m => ZFunctionDef -> m ()
addFunctionDef def = def `seq` modify $ \z -> z 
  { zFunctions = IMap.insert (zVarId $ funcName def) def (zFunctions z) }
  
addLemma :: MonadState Zeno m => String -> ZClause -> m ()
addLemma name lemma = lemma `seq` 
  modify $ \z -> z { zLemmas = Map.insert name lemma (zLemmas z) }
  
showZId :: MonadReader Zeno m => ZId -> m String
showZId id = (fromMaybe (showId id)) `liftM` (lookupName id)

lookupName :: MonadReader Zeno m => ZId -> m (Maybe String)
lookupName id = Map.lookup id `liftM` zNames `liftM` ask 

addGlobalVar :: MonadState Zeno m => ZVar -> String -> m ()
addGlobalVar var name = do 
  christen (zVarId var) name
  modify $ \z -> z 
    { zGlobalVars = Map.insert name var (zGlobalVars z) }
    
getGlobalVar :: MonadReader Zeno m => String -> m ZVar
getGlobalVar name = do
  vars <- zGlobalVars `liftM` ask
  case Map.lookup name vars of
    Nothing -> error $ "Global variable \"" ++ name ++ "\" not found"
    Just gv -> return gv
{-
findFunction :: MonadReader Zeno m => String -> m ZFunctionDef
findFunction name = do
  funs <- zFunctions `liftM` ask
  funs' <- mapKeysM showZVar funs 
  case Map.lookup name funs' of
    Nothing -> error $ "Function \"" ++ name ++ "\" not found"
    Just fn -> return fn
  -}  
findDataType :: MonadReader Zeno m => String -> m ZDataType
findDataType name = do
  dtypes <- zDataTypes `liftM` ask
  names <- mapM (showZId . zDataTypeId) dtypes
  case lookup name (names `zip` dtypes) of
    Nothing -> error $ "Datatype \"" ++ name ++ "\" not found"
    Just dt -> return dt
    
instantiate :: ZVar -> ZenoState (ZTerm, [ZVar])
instantiate var = do
  let var_types = flattenType (zVarType var)
      (arg_types, [res_type]) = splitAt (length var_types - 1) var_types
  new_ids <- replicateM (length arg_types) newIdS
  let args' = zipWith ($) (map ZVar new_ids) arg_types
      args = map ($ ZVariable) args'
      vars = filter ((== res_type) . zVarType) args
      term = unflattenTerm (map Var (var : args))
  return (term, vars)
  
applyTermDefinitions :: TermTraversable t => 
    [ZEquality] -> t ZVar -> ZenoReader ((t ZVar), Bool)
applyTermDefinitions defs t = do
  let mapped_t = mapTermsM applyTermDefinition t
  (t', any) <- runWriterT mapped_t
  if getAny any 
    then do
      (t'', _) <- applyTermDefinitions defs t'
      return (t'', True)
    else 
      return (t', False)
  where
    applyTermDefinition :: ZTerm -> WriterT Any ZenoReader ZTerm
    applyTermDefinition term = do
      expr <- lift $ applyFunctionDef defs term
      case expr of 
        Just (Value term') -> tell (Any True) >> return term'
        _ -> return term
    
getName :: MonadReader Zeno m => ZId -> m String
getName = (liftM fromJust) . lookupName

largestId :: ZClause -> ZId
largestId = maximum . (map zVarId) . toList

showZTerm :: MonadReader Zeno m => ZTerm -> m String
showZTerm = liftM showUnquoted . mapM showZVar

showZVar :: MonadReader Zeno m => ZVar -> m String
showZVar = showZId . zVarId

showZVarTyped :: MonadReader Zeno m => ZVar -> m String
showZVarTyped var = do
  var_s <- showZVar var
  type_s <- showZType (zVarType var)
  return (var_s ++ " : " ++ type_s)

showZType :: MonadReader Zeno m => ZType -> m String
showZType = liftM showUnquoted . mapM showZDataType

showZDataType :: MonadReader Zeno m => ZDataType -> m String
showZDataType = showZId . zDataTypeId
 
showFullZDataType :: MonadReader Zeno m => ZDataType -> m String
showFullZDataType datatype = do
  name_s <- showZDataType datatype
  cons_s <- mapM showZVarTyped (zDataTypeCons datatype)
  let rec = if recursiveDataType datatype then "rec " else ""
  return $ rec ++ "type " ++ name_s ++ "\n  " 
    ++ (intercalate "\n  " cons_s)
     
showZEquality :: MonadReader Zeno m => ZEquality -> m String
showZEquality (t1, t2) = do
  t1' <- showZTerm t1
  t2' <- showZTerm t2
  return $ showEquality (t1', t2')
    
showZClause :: MonadReader Zeno m => ZClause -> m String
showZClause = (liftM showUnquoted) . mapM showZVar

showZTermSub :: MonadReader Zeno m => ZTermSubstitution -> m String
showZTermSub = (liftM show) . mapM (mapPairM showZTerm) . toList

showZFunctionDef :: MonadReader Zeno m => ZFunctionDef -> m String
showZFunctionDef = (liftM showUnquoted) . mapM showZVar

showZExpression :: MonadReader Zeno m => ZExpression -> m String
showZExpression = (liftM showUnquoted) . mapM showZVar

showZProof :: ZProof -> ZenoReader String
showZProof proof = do
  let inds = findInducts True proof
  proof' <- mapM (fmap showUnquoted . mapM showZVar) proof
  inds_s <- intercalate "\n" <$> mapM showZClause inds
  return $ showUnquoted proof' ++
    "\n\nLemmas Proven:\n\n" ++ inds_s
  where
    findInducts :: Bool -> ZProof -> [ZClause]
    findInducts _ Unknown = []
    findInducts _ (Disproven {}) = []
    findInducts ind (Proven (ProofStep step cls) ps) =
      let ind_step = "ind" `isPrefixOf` step
          ind' = not ind_step && ("hyp" `isPrefixOf` step || ind)
          rest = concatMap (findInducts ind') ps
      in if ind_step && ind then cls : rest else rest
      

instance Show ZVar where
  show = show . zVarId

instance Show Zeno where
  show zeno = runReader reader zeno
    where
      datatypes = zDataTypes zeno
      functions = IMap.elems $ zFunctions zeno
      names = Map.toList $ zNames zeno
      lemmas = Map.toList $ zLemmas zeno
      
      reader :: MonadReader Zeno m => m String
      reader = do
        datatypes' <- mapM showFullZDataType datatypes
        let datatypes_s = intercalate "\n\n" datatypes'
        
        functions' <- mapM (mapM showZVar) functions
        let functions_s = intercalate "\n\n" (map showUnquoted functions')
        
        let showLemma (name, lemma) = do
              lemma' <- showZClause lemma
              return $ name ++ "\n  " ++ lemma'
        lemmas' <- mapM showLemma lemmas
        let lemmas_s = intercalate "\n\n" lemmas'
        
        return $ "\n\n" ++ datatypes_s ++ 
          "\n\n" ++ functions_s ++ 
          "\n\n" ++ lemmas_s

        

